import numpy as np
import matplotlib.pyplot as plt
import math
import mpl_toolkits.axisartist as axisartist


def func1(t, T, delta):
    # 方波信号
    if (t + delta) % T <= 2 * delta:
        return 1
    else:
        return 0

def func2(f, delta):
    # 方波信号的双边谱
    return math.sin(2 * math.pi * f * delta) / (math.pi * f)

def set_axis(fig, number, title, x_range=(-10, 15), y_range=(-1, 1)):
    ax = axisartist.Subplot(fig, number)
    fig.add_axes(ax)
    ax.axis[:].set_visible(False)
    ax.axis["x"] = ax.new_floating_axis(0,0)
    ax.axis["x"].set_axisline_style("->", size = 1.0)
    ax.axis["y"] = ax.new_floating_axis(1,0)
    ax.axis["y"].set_axisline_style("->", size = 1.0)
    ax.axis["x"].set_axis_direction("bottom")
    ax.axis["y"].set_axis_direction("left")
    ax.set_xlim(x_range[0], x_range[1])
    ax.set_ylim(y_range[0], y_range[1])
    ax.set_title(title)

def mod(x, f):
    x = abs(x)
    return x - math.floor(x / f) * f < 1e-2


if __name__ == "__main__":

    fig = plt.figure(figsize=(12, 8))
    plt.rcParams['font.family'] = ['Arial Unicode MS']

    # 绘制原方波信号
    T = 3
    delta = 1
    low = -T * 1.1
    high = T * 1.1
    t = np.arange(-2 * T, 2 * T, 0.01)

    f1 = np.array([func1(i, T, delta) for i in t])
    set_axis(fig, 421, "周期为$T_0$的方波信号", (low, high), (-1, 2))
    plt.plot(t, f1)

    # 绘制周期扩大的方波信号
    T = 12
    delta = 1
    low = -T * 1.1
    high = T * 1.1
    t = np.arange(-2 * T, 2 * T, 0.01)

    f2 = np.array([func1(i, T, delta) for i in t])
    set_axis(fig, 423, "周期为$4T_0$的方波信号", (low, high), (-1, 2))
    plt.plot(t, f2)

    # 绘制新的方波信号频谱
    f = np.arange(-2 * T, 2 * T, 0.01)
    n_f = np.array([i for i in f if mod(i, 1 / T)])
    
    f3 = np.array([1 / T * func2(i, delta) for i in n_f])
    set_axis(fig, 424, "新方波信号的频谱", (-2, 2), (min(f3) * 1.2, max(f3) * 1.2))
    plt.stem(n_f, f3, markerfmt='.')

    # 绘制周期16T的方波信号
    T = 48
    delta = 1
    low = -T * 1.1
    high = T * 1.1
    t = np.arange(-2 * T, 2 * T, 0.01)

    f4 = np.array([func1(i, T, delta) for i in t])
    set_axis(fig, 425, "周期为$16T_0$的方波信号", (low, high), (-1, 2))
    plt.plot(t, f4)

    # 绘制16T方波信号频谱
    f = np.arange(-2 * T, 2 * T, 0.01)
    n_f = np.array([i for i in f if mod(i, 1 / T)])
    
    f5 = np.array([1 / T * func2(i, delta) for i in n_f])
    set_axis(fig, 426, "周期为$16T_0$的方波信号的频谱", (-1.5, 1.5), (min(f5) * 1.2, max(f5) * 1.2))
    plt.stem(n_f, f5, markerfmt='.')

    # 绘制周期64T的方波信号
    T = 3 * 64
    delta = 2
    low = -T * 1.1
    high = T * 1.1
    t = np.arange(-2 * T, 2 * T, 0.01)

    f6 = np.array([func1(i, T, delta) for i in t])
    set_axis(fig, 427, "周期为$64T_0$的方波信号", (low, high), (-1, 2))
    plt.plot(t, f6)

    # 绘制16T方波信号的频谱
    f = np.arange(-2 * T, 2 * T, 0.01)
    n_f = np.array([i for i in f if mod(i, 1 / T)])
    
    f7 = np.array([1 / T * func2(i, delta) for i in n_f])
    set_axis(fig, 428, "周期为$64T_0$的方波信号的频谱", (-0.75, 0.75), (min(f7) * 1.2, max(f7) * 1.2))
    plt.stem(n_f, f7, markerfmt='.')
    
    plt.show()